package interceptors

import (
	"time"

	"github.com/gin-gonic/gin"
	"github.com/juju/ratelimit"
	"github.com/spf13/viper"
)

// RateLimit 限流中间件
func RateLimit() gin.HandlerFunc {
	bucket := ratelimit.NewBucket( //ms
		time.Duration(viper.GetInt("ratelimit.fillInterval"))*time.Microsecond,
		viper.GetInt64("ratelimit.capacity"),
	)
	return func(c *gin.Context) {
		bucket.Wait(1)
		c.Next()
	}
}
